# -*- encoding: utf-8 -*-
import os
import os.path
from BeautifulSoup import BeautifulSoup
from elasticsearch import Elasticsearch
import click
import logging
import time
import sys


logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S',
                    filename=os.path.dirname(os.path.dirname(sys.path[0]))+'/log/'+time.strftime('%Y-%m-%d', time.localtime(time.time())) + '.log',
                    filemode='w')


@click.command()
@click.option('--filepath', help='html file path', required=1)
@click.option('--host', help='elasticsearch host address', default='http://127.0.0.1:9200/',  required=1)
def main(filepath, host):
    egg_file(filepath, host)


def egg_file(filepath, host):
    for parent,dirnames,filenames in os.walk(filepath):
        for filename in filenames:
            if filename.endswith('.html'):
                logging.info("filename is: %s", filename)
                filename_all_path = filepath + filename
                dealWithHtml(filename, filename_all_path, host)


def insert_htmldb_all(es, htmlid, VulnDBID, LastModified, PublishedDate, title, timeLine_dict, vtemMetrics_dict,
                      description, soultion, classification_dict, keywords, affectedproducts_lst,
                      affectedCpe_lst, references_dict, cvsSv2Score_lst, credit_lst, cve_id, abstract, cwe_id,
                      impact_dict, cve_reference_dict, cve_references_lst, vulnerable_configurations):
    obj = {
        'htmlid': htmlid,
        'VulnDBID': VulnDBID,
        'LastModified': LastModified,
        'PublishedDate': PublishedDate,
        'title': title,
        'timeLine_dict': timeLine_dict,
        'vtemMetrics_dict': vtemMetrics_dict,
        'description': description,
        'soultion': soultion,
        'classification_dict': classification_dict,
        'keywords': keywords,
        'affectedproducts_lst': affectedproducts_lst,
        'affectedCpe_lst': affectedCpe_lst,
        'references_dict': references_dict,
        'cvsSv2Score_lst': cvsSv2Score_lst,
        'credit_lst': credit_lst,
        'cve_id': cve_id,
        'abstract': abstract,
        'cwe_id': cwe_id,
        'impact_dict': impact_dict,
        'cve_reference_dict': cve_reference_dict,
        'cve_references_lst': cve_references_lst,
        'vulnerable_configurations': vulnerable_configurations,
    }
    return es.index(index="htmldb", doc_type="html", body=obj, id=htmlid, refresh=True)


def dealWithHtml(filename, filename_all_path, host):
    f = open(filename_all_path, "r")
    html_str = f.read()
    f.close()

    htmles = Elasticsearch(hosts=host)
    htmles.indices.create(index='htmldb', ignore=400)
    if not htmles.indices.get_mapping(index='htmldb', doc_type='html'):
        htmles.indices.put_mapping(index='htmldb', doc_type='html', body={
            "_all": {"enabled": True},
            "properties": {
                "html_id": {"type": "string", "index": "analyzed"},
            }
        })

    htmlid = filename.split('.')[0]
    soup = BeautifulSoup(html_str)
    metadata = soup.find(id="vuln-metadata")
    li1 = metadata.contents[0]
    VulnDBID = li1.text
    li2 = metadata.contents[1]
    LastModified = li2.text
    li3 = metadata.contents[2]
    PublishedDate = li3.text
    title = metadata.nextSibling
    title_text = title.text
    timeLine = title.nextSibling
    dateAndValue = timeLine.contents[1]
    timeLine_dict = {}
    for i in range(0, len(dateAndValue.contents[0])):
        dateName_ori = dateAndValue.contents[0].contents[i].text
        dateName = dateName_ori.replace(' ', '_')
        dateValue = dateAndValue.contents[1].contents[i].text
        if dateValue == 'N/A':
            dateValue = None
        timeLine_dict[dateName] = dateValue

    VTEMMetrics = timeLine.nextSibling
    timeAndValue = VTEMMetrics.contents[1]
    vtemMetrics_dict = {}
    for i in range(0, len(timeAndValue.contents[0])):
        timeName_ori = timeAndValue.contents[0].contents[i].text
        timeName = timeName_ori.replace(' ', '_')
        timeValue = timeAndValue.contents[1].contents[i].text
        if timeValue == 'N/A':
            timeValue = None
        vtemMetrics_dict[timeName] = timeValue

    description = VTEMMetrics.nextSibling
    descriptionValue = description.nextSibling.text

    soultion = description.nextSibling.nextSibling
    soultionValue = soultion.nextSibling.text

    classification = soultion.nextSibling.nextSibling
    nameAndValue = classification.contents[1]
    classification_dict = {}
    for i in range(0, len(nameAndValue.contents[0])):
        name_ori = nameAndValue.contents[0].contents[i].text
        name = name_ori.replace(' ', '_')
        value = nameAndValue.contents[1].contents[i].text
        classification_dict[name] = value

    global keywordsValue
    keywords = classification.nextSibling
    keywords_text = keywords.text
    keywordsValue = ''
    if keywords_text == 'Keywords':
        keywordsValue = keywords.nextSibling.text
        AffectedProducts = keywords.nextSibling.nextSibling
    else:
        AffectedProducts = classification.nextSibling
    infoAndValue = AffectedProducts.contents[0].contents[1]
    affectedproducts_lst = []
    for i in range(1, len(infoAndValue.contents)):
        global vendor, product, version
        if len(infoAndValue.contents[i].contents) == 3:
            vendor = infoAndValue.contents[i].contents[0].text
            product = infoAndValue.contents[i].contents[1].text
            version = infoAndValue.contents[i].contents[2].text
        elif len(infoAndValue.contents[i].contents) == 2:
            product = infoAndValue.contents[i].contents[0].text
            version = infoAndValue.contents[i].contents[1].text
        elif len(infoAndValue.contents[i].contents) == 1:
            version = infoAndValue.contents[i].contents[0].text
        else:
            continue
        affectedproducts_lst.append(vendor + ':' + product + ':' + version)

    AffectedCPEs = AffectedProducts.nextSibling
    AffectedCPEs_text = AffectedCPEs.contents[0].contents[0].text
    affectedCpe_lst = []
    if AffectedCPEs_text == 'Affected CPEs':
        for i in range(0, len(AffectedCPEs.contents[0].contents[1].contents[0].contents)):
            cpe = AffectedCPEs.contents[0].contents[1].contents[0].contents[i].text
            affectedCpe_lst.append(cpe)
        references = AffectedCPEs.nextSibling.nextSibling
    else:
        references = AffectedProducts.nextSibling.nextSibling

    references_dict = {}
    for i in range(1, len(references.contents)):
        referencesKey_ori = references.contents[i].contents[0].text
        referencesKey = referencesKey_ori.replace(' ', '_')
        referencesValue = references.contents[i].contents[1].text
        references_dict[referencesKey] = referencesValue

    temp = references
    cvsSv2Score_lst = []
    global credit
    if temp.nextSibling.text != 'CVSSv2 Score':  # 不存在 CVSSv2 Score 项时，credit 的游标位置
        credit = temp.nextSibling
    else:
        while True:
            CVSSv2Score_temp = temp.nextSibling.nextSibling
            try:
                SourceGeneratedScore = CVSSv2Score_temp.text
            except:
                SourceGeneratedScore = CVSSv2Score_temp
            if not SourceGeneratedScore.startswith('Source'):
                break
            cvsSv2Score_dict = {}
            SourceGeneratedScore_splt = SourceGeneratedScore.split('|')
            cvsSv2Score_dict['Source'] = SourceGeneratedScore_splt[0].split(':')[1]
            cvsSv2Score_dict['Generated'] = SourceGeneratedScore_splt[1].split(':')[1]
            cvsSv2Score_dict['Score_Base_Score'] = SourceGeneratedScore_splt[2].split('=')[1]
            CVSSv2Score = CVSSv2Score_temp

            relateScore = CVSSv2Score.nextSibling
            temp = CVSSv2Score
            for i in range(0, len(relateScore.contents[0].contents)):
                relateScoreKey_ori = relateScore.contents[0].contents[i].text
                relateScoreKey = relateScoreKey_ori.replace(' ', '_')
                relateScoreValue = relateScore.contents[1].contents[i].text
                cvsSv2Score_dict[relateScoreKey] = relateScoreValue
            cvsSv2Score_lst.append(cvsSv2Score_dict)
            credit = relateScore.nextSibling

    if credit.contents[0].text == 'Change Log':
        credit = credit.nextSibling

    credit_lst = []
    if credit.contents[0].text != 'Credit':
        global CVE
        CVE = credit
    else:
        for i in range(0, len(credit.contents[1].contents)):
            creditValue = credit.contents[1].contents[i].text
            credit_lst.append(creditValue)
        CVE = credit.nextSibling

    if CVE == None:
        logging.info("insert all data to htmldb with CVE is None")
        insert_htmldb_all(htmles, htmlid, VulnDBID, LastModified, PublishedDate, title_text, timeLine_dict,
                          vtemMetrics_dict, descriptionValue, soultionValue, classification_dict, keywords_text,
                          affectedproducts_lst, affectedCpe_lst, references_dict, cvsSv2Score_lst, credit_lst,
                          cve_id=None, abstract=None, cwe_id=None, impact_dict=None, cve_reference_dict=None,
                          cve_references_lst=None, vulnerable_configurations=None)
        return
    root = CVE.contents[0].contents[0]
    global cve_id
    cve_id = ''
    try:
        cve_id = root.contents[0].contents[1].text
    except:
        cve_id = root.contents[0].contents[1]
    abstract = root.contents[1].contents[0].text

    cwe_id = ''
    impact_dict = {}
    cve_reference_dict = {}
    cve_references_lst = []
    vulnerable_configurations = ''
    for i in range(01, len(root.contents[1].contents)):
        try:
            tag_value = root.contents[1].contents[i].text
        except:
            tag_value = ''
        if tag_value == 'CWE-ID':
            cwe_id = root.contents[1].contents[i+1].text
        elif tag_value == 'Score':
            score = root.contents[1].contents[i+1]
            impact_dict['Score'] = score
        elif tag_value == 'Access vector':
            acc_vct = root.contents[1].contents[i+1]
            impact_dict['Access_vector'] = acc_vct
        elif tag_value == 'Access complexity':
            acc_cmp = root.contents[1].contents[i+1]
            impact_dict['Access_complexity'] = acc_cmp
        elif tag_value == 'Authentication':
            ath = root.contents[1].contents[i+1]
            impact_dict['Authentication'] = ath
        elif tag_value == 'Confidentiality impact':
            cfd_imp = root.contents[1].contents[i+1]
            impact_dict['Confidentiality_impact'] = cfd_imp
        elif tag_value == 'Integrity impact':
            itg_imp = root.contents[1].contents[i+1]
            impact_dict['Integrity_impact'] = itg_imp
        elif tag_value == 'Availability impact':
            avai_imp = root.contents[1].contents[i+1]
            impact_dict['Availability_impact'] = avai_imp
        elif tag_value == 'Source':
            source = root.contents[1].contents[i+1]
            cve_reference_dict['source'] = source
        elif tag_value == 'name':
            name = root.contents[1].contents[i+1]
            cve_reference_dict['name'] = name
        elif tag_value == 'link':
            link = root.contents[1].contents[i+1]
            cve_reference_dict['link'] = link
            cve_references_lst.append(cve_reference_dict)
        elif tag_value == 'Vulnerable Configurations':
            vln_conf = ''
            for j in range(i+1, len(root.contents[1].contents)):
                try:
                    temp_str = root.contents[1].contents[j].text
                except:
                    temp_str = root.contents[1].contents[j]
                if temp_str:
                    vln_conf = vln_conf + ' ' + temp_str
            vulnerable_configurations = vln_conf
    logging.info("insert all data to htmldb with CVE is not None")
    insert_htmldb_all(htmles, htmlid, VulnDBID, LastModified, PublishedDate, title_text, timeLine_dict,
                      vtemMetrics_dict, descriptionValue, soultionValue, classification_dict, keywords_text,
                      affectedproducts_lst, affectedCpe_lst, references_dict, cvsSv2Score_lst, credit_lst,
                      cve_id, abstract, cwe_id, impact_dict, cve_reference_dict, cve_references_lst,
                      vulnerable_configurations)

if __name__ == '__main__':
    main()
